import pandas as pd
import matplotlib.pyplot as plt
import os
import multiprocessing


def plot_navs(frequency, file):
    """绘制净值曲线"""

    print(file)
    nav_df = pd.read_csv(f'{frequency}/nav/{file}')
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))

    # 净值曲线
    ax1.plot(nav_df['date'], nav_df['nav0'], label=f'nav', color='purple')
    ax1.plot(nav_df['date'], nav_df['nav2'], label=f'nav_after_fee', color='blue')
    ax1.plot(nav_df['date'], nav_df['nav(hold_btc)'], label='Hold BTC', color='red')
    ax1.set_title('Strategy Performance')
    ax1.set_ylabel('NAV')
    ax1.set_xlabel('Date')
    ax1.legend()

    ax2.plot(nav_df['date'], nav_df['nav0'], label=f'nav', color='purple')
    ax2.plot(nav_df['date'], nav_df['nav2'], label=f'nav_after_fee', color='blue')
    ax2.plot(nav_df['date'], nav_df['nav(hold_btc)'], label='Hold BTC', color='red')
    ax2.set_yscale('log')  # 设置 y 轴为对数尺度
    ax2.set_title('Strategy Performance (Logarithmic Scale)')
    ax2.set_ylabel('NAV(log)')
    ax2.set_xlabel('Date')
    ax2.legend()

    plt.tight_layout()
    plt.savefig(f'{frequency}/png/{file[:-4]}.png')
    plt.close(fig)  # 每次循环后关闭


if __name__ == '__main__':
    # 画全部净值图
    for frequency in ["1d"]:
        files = os.listdir(f'./{frequency}/nav')
        params = []
        for file in files:
            params.append((frequency, file))
        with multiprocessing.Pool(60) as pool:
            pool.starmap(plot_navs, params)